import os
import numpy as np
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score
from sklearn.metrics import davies_bouldin_score
from sklearn.metrics import calinski_harabasz_score
import matplotlib.pyplot as plt

"""
silhouette_score Method for Optimal k
"""

max_label_num = 50000


def read_labels(label_file):
    objects = []
    with open(label_file, "r") as f:
        lines = f.readlines()
    for line in lines:
        parts = line.strip().split()
        if len(parts) > 1:
            width = float(parts[3]) - float(parts[1])
            height = float(parts[4]) - float(parts[2])
            objects.append((width, height))
    return objects


def calculate_aspect_ratio(objects):
    aspect_ratios = [w / h for w, h in objects if h != 0]
    return aspect_ratios


def plot_metrics(aspect_ratios):
    silhouette_scores = []
    db_scores = []
    ch_scores = []
    sse_scores = []
    K_range = range(2, 6)
    for k in K_range:
        if len(aspect_ratios) < k:
            continue

        kmeans = KMeans(n_clusters=k, random_state=0)
        cluster_labels = kmeans.fit_predict(aspect_ratios)

        # Calculating metrics
        silhouette_avg = silhouette_score(aspect_ratios, cluster_labels)
        db_index = davies_bouldin_score(aspect_ratios, cluster_labels)
        ch_index = calinski_harabasz_score(aspect_ratios, cluster_labels)
        sse = kmeans.inertia_  # SSE score
        print(
            f"k:{k}\nsilhouette:{silhouette_avg}\ndb_index:{db_index}\nch_index:{ch_index}\nsse:{sse}\n"
        )
        # Storing scores
        silhouette_scores.append(silhouette_avg)
        db_scores.append(db_index)
        ch_scores.append(ch_index)
        sse_scores.append(sse)

    # Plotting all scores
    fig, axs = plt.subplots(
        2, 2, figsize=(15, 10)
    )  # Create a figure and a grid of subplots

    # Plot Silhouette Scores
    axs[0, 0].plot(
        K_range,
        silhouette_scores,
        "o-",
        color="blue",
        marker="o",
        markersize=8,
        linewidth=2,
    )
    axs[0, 0].set_title("Silhouette Score", fontsize=14)
    axs[0, 0].set_xlabel("Number of Clusters (k)")
    axs[0, 0].set_ylabel("Silhouette Score")
    # General layout adjustments
    plt.tight_layout()
    plt.grid(True, linestyle="--", alpha=0.6)

    # Plot Davies-Bouldin Index
    axs[0, 1].plot(
        K_range, db_scores, "o-", color="red", marker="o", markersize=8, linewidth=2
    )
    axs[0, 1].set_title("Davies-Bouldin Index", fontsize=14)
    axs[0, 1].set_xlabel("Number of Clusters (k)")
    axs[0, 1].set_ylabel("Davies-Bouldin Index")
    # General layout adjustments
    plt.tight_layout()
    plt.grid(True, linestyle="--", alpha=0.6)

    # Plot Calinski-Harabasz Index
    axs[1, 0].plot(
        K_range, ch_scores, "o-", color="green", marker="o", markersize=8, linewidth=2
    )
    axs[1, 0].set_title("Calinski-Harabasz Index", fontsize=14)
    axs[1, 0].set_xlabel("Number of Clusters (k)")
    axs[1, 0].set_ylabel("Calinski-Harabasz Index")

    # General layout adjustments
    plt.tight_layout()
    plt.grid(True, linestyle="--", alpha=0.6)

    # Plot SSE
    axs[1, 1].plot(
        K_range, sse_scores, "o-", color="purple", marker="o", markersize=8, linewidth=2
    )
    axs[1, 1].set_title("SSE (Sum of Squared Errors)", fontsize=14)
    axs[1, 1].set_xlabel("Number of Clusters (k)")
    axs[1, 1].set_ylabel("SSE")

    # General layout adjustments
    plt.tight_layout()
    plt.grid(True, linestyle="--", alpha=0.6)
    # 使用plt.savefig()并指定参数
    plt.savefig("my_figure.jpg", dpi=1000, bbox_inches="tight", pad_inches=0)

    plt.show()


# Assuming the function call and data preprocessing is handled in main or similar function
def main():
    folder_path = "/home/hw/dataset/cone_total_new/labels"
    label_files = [f for f in os.listdir(folder_path) if f.endswith(".txt")]

    aspect_ratios = []
    print("文件数量", len(label_files))
    print("计算中")
    for label_file in label_files:
        objects = read_labels(os.path.join(folder_path, label_file))
        aspect_ratios.extend(calculate_aspect_ratio(objects))

    aspect_ratios = np.array(aspect_ratios).reshape(-1, 1)
    print("目标数量", len(aspect_ratios))
    if len(aspect_ratios) > max_label_num:
        aspect_ratios = aspect_ratios[:max_label_num]
    print("最后目标数量", len(aspect_ratios))
    plot_metrics(aspect_ratios)


if __name__ == "__main__":
    main()
